#!/usr/bin/env python3
import math
import cv2
import apriltag
import numpy as np
import cv_bridge
import rclpy
from rclpy.node import Node
from std_msgs.msg import Int32
from sensor_msgs.msg import Image
from sensor_msgs.msg import CameraInfo
from tf2_ros.transform_broadcaster import TransformBroadcaster
from geometry_msgs.msg import TransformStamped
from rcl_interfaces.msg import ParameterDescriptor
import tf_transformations as tft


class AprilTagFinder(Node):
    def __init__(self):
        super().__init__("april_tag_finder")
        self.declare_parameter(
            "tag_family",
            "tag36h11",
            ParameterDescriptor(description="which apriltag tag family used."),
        )
        self.declare_parameter(
            "tag_size",
            1.0,
            ParameterDescriptor(
                description="the size of apriltag tag in meter(need more check)."
            ),
        )
        self.declare_parameter(
            "debug",
            False,
            ParameterDescriptor(
                description="whether to show debug info and creat postprocess image."
            ),
        )
        self.declare_parameter("use_ns_as_tf_prefix", True)
        tag_family = self.get_parameter("tag_family").get_parameter_value().string_value
        self.tag_size = (
            self.get_parameter("tag_size").get_parameter_value().double_value
        )
        self.debug: bool = self.get_parameter("debug").get_parameter_value().bool_value
        self.useNS: bool = (
            self.get_parameter("use_ns_as_tf_prefix").get_parameter_value().bool_value
        )
        self.targetIDPub = self.create_publisher(Int32, "target_id", 1)
        self.publisher = self.create_publisher(Image, "image_processed", 1)
        self.subscription = self.create_subscription(
            Image, "camera_image", self.image_callback, 1
        )
        self.cameraInfoSub = self.create_subscription(
            CameraInfo, "camera_info", self.camera_info_callback, 1
        )
        self.cam_params = None
        self.tf_broadcaster = TransformBroadcaster(self)

        self.tag_detector = apriltag.Detector(
            apriltag.DetectorOptions(families=tag_family)
        )

        self.parent_frame = self.get_namespace() + "/camera" if self.useNS else "camera"
        self.child_frame = self.get_namespace() + "/tag" if self.useNS else "tag"

        self.lastRotation = None

    def image_callback(self, msg: Image):
        cv_image = cv_bridge.CvBridge().imgmsg_to_cv2(msg, desired_encoding="bgr8")
        gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)
        tags = self.tag_detector.detect(gray)
        for tag in tags:
            # self.targetIDPub.publish(Int32(tag.id))
            M, e1, e2 = self.tag_detector.detection_pose(
                tag, self.cam_params, self.tag_size
            )
            transpose = M[0:3, 3]
            rotation = M[0:3, 0:3]
            # self.get_logger().debug(
            #    f"x:{transpose[0]}, y:{transpose[1]}, z:{transpose[2]}"
            # )
            t = TransformStamped()
            t.header.stamp = self.get_clock().now().to_msg()
            t.header.frame_id = self.parent_frame
            # t.child_frame_id = "tag" + str(tag.id)
            t.child_frame_id = self.child_frame
            t.transform.translation.x = -transpose[0]
            t.transform.translation.y = -transpose[1]
            t.transform.translation.z = transpose[2]
            q = tft.quaternion_from_matrix(M.T)

            if self.lastRotation is None:
                self.lastRotation = q
            self.lastRotation = tft.quaternion_slerp(self.lastRotation, q, 0.5)
            q = self.lastRotation

            t.transform.rotation.x = q[0]
            t.transform.rotation.y = q[1]
            t.transform.rotation.z = q[2]
            t.transform.rotation.w = q[3]
            self.tf_broadcaster.sendTransform(t)
            # print(f"{self.parent_frame}:{self.child_frame}")

            if self.debug:
                for i in range(4):
                    cv2.circle(
                        cv_image, tuple(tag.corners[i].astype(int)), 4, (255, 0, 0), 2
                    )
                cv2.circle(cv_image, tuple(tag.center.astype(int)), 4, (2, 180, 200), 4)
            # current only want one tag be published
            break
        if self.debug:
            self.publisher.publish(cv_bridge.CvBridge().cv2_to_imgmsg(cv_image, "bgr8"))

    def camera_info_callback(self, msg: CameraInfo):
        fx = msg.k[2]  # [0,2]
        fy = msg.k[5]  # [1,2]
        cx = msg.k[0]  # [0,0]
        cy = msg.k[4]  # [1,1]
        self.cam_params = [cx, cy, fx, fy]


def main(args=None):
    rclpy.init(args=args)
    node = AprilTagFinder()
    node.get_logger().info("AprilTag Finder node started")
    rclpy.spin(node)
    rclpy.shutdown()


if __name__ == "__main__":
    main()
